import cv2
import os
import argparse
import logging
from pathlib import Path
from datetime import datetime

def setup_logger(output_dir):
    log_dir = os.path.join(output_dir, 'logs')
    os.makedirs(log_dir, exist_ok=True)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(log_dir, f"capture_{timestamp}.log")
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )
    
    return logging.getLogger()

def extract_frames(video_path, output_dir, interval, start_time=0, end_time=None, millisecond_digits=6):
    logger = setup_logger(output_dir)
    
    logger.info(f"开始处理视频: {video_path}")
    logger.info(f"输出目录: {output_dir}")
    logger.info(f"截图间隔: {interval}秒")
    logger.info(f"开始时间: {start_time}秒")
    logger.info(f"结束时间: {end_time if end_time is not None else '视频结束'}")
    logger.info(f"毫秒位数: {millisecond_digits}")
    
    os.makedirs(output_dir, exist_ok=True)
    
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        logger.error(f"无法打开视频文件: {video_path}")
        return
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count / fps
    
    video_name = Path(video_path).stem
    
    start_time = max(0, start_time)
    if end_time is None:
        end_time = duration
    end_time = min(end_time, duration)
    
    if start_time >= end_time:
        logger.error(f"开始时间 {start_time}秒 大于或等于结束时间 {end_time}秒")
        cap.release()
        return
    
    logger.info(f"视频总时长: {duration:.2f}秒")
    logger.info(f"有效截图区间: {start_time:.2f}秒 到 {end_time:.2f}秒")
    
    # 计算需要截取的总帧数
    expected_frames = int((end_time - start_time) / interval) + 1
    logger.info(f"预计截取帧数: {expected_frames}")
    
    # 设置到开始帧
    cap.set(cv2.CAP_PROP_POS_MSEC, start_time * 1000)
    
    start_process_time = datetime.now()
    logger.info(f"开始截图处理...")
    
    current_time = start_time
    image_count = 0
    
    # 循环直到超出结束时间
    while current_time <= end_time + interval/2:  # 增加半间隔容差
        # 设置到当前时间点
        cap.set(cv2.CAP_PROP_POS_MSEC, current_time * 1000)
        
        # 读取当前帧
        success, image = cap.read()
        if not success:
            logger.warning(f"无法读取帧 @ {current_time:.3f}秒")
            break
        
        # 计算当前时间的毫秒表示
        current_millisecond = int(current_time * 1000)
        
        # 构建输出文件名
        millisecond_str = str(current_millisecond)
        if len(millisecond_str) <= millisecond_digits:
            millisecond_str = millisecond_str.zfill(millisecond_digits)
        output_filename = f"{video_name}_{millisecond_str}.jpg"
        output_path = os.path.join(output_dir, output_filename)
        
        # 保存图像
        try:
            cv2.imwrite(output_path, image)
            image_count += 1
            logger.info(f"已保存: {output_path} (时间点: {current_time:.3f}秒)")
        except Exception as e:
            logger.error(f"保存图像失败 @ {current_time:.3f}秒: {str(e)}")
        
        # 更新当前时间
        current_time += interval
        
        # 更新进度
        if image_count % 10 == 0 or current_time >= end_time:
            progress = min(100.0, (current_time - start_time) / (end_time - start_time) * 100)
            logger.info(f"进度: {progress:.1f}% ({image_count}/{expected_frames}帧)")
    
    end_process_time = datetime.now()
    process_duration = (end_process_time - start_process_time).total_seconds()
    
    cap.release()
    
    logger.info(f"截图处理完成!")
    logger.info(f"总耗时: {process_duration:.2f}秒")
    logger.info(f"共截取: {image_count}张图像")
    logger.info(f"日志文件位置: {os.path.join(output_dir, 'logs')}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='从视频中按指定间隔截取图像')
    parser.add_argument('video_path', help='输入视频文件的路径')
    parser.add_argument('output_dir', help='输出图像的目录')
    parser.add_argument('--interval', type=float, default=1.0, help='截图间隔（秒），默认值为1.0')
    parser.add_argument('--start', type=float, default=0, help='开始时间（秒），默认值为0')
    parser.add_argument('--end', type=float, default=None, help='结束时间（秒），默认值为视频结束')
    parser.add_argument('--millis-digits', type=int, default=6, help='毫秒位数，默认值为6')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.video_path):
        print(f"错误：视频文件 {args.video_path} 不存在")
    else:
        extract_frames(args.video_path, args.output_dir, args.interval, args.start, args.end, args.millis_digits)